from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageChops, ImageFile
from io import BytesIO
import os
import math
ImageFile.LOAD_TRUNCATED_IMAGES = True
Image.MAX_IMAGE_PIXELS = None

class WaterMark:
    """WaterMark
        source: 图片的字节流
        text: 添加的水印文本
        color: 表示添加的颜色
        space: 水印间隔
        angle: 水印角度
        font: 字体类型
        height: 字体高度
        txtsize: 字体大小
        opacity: 透明度
        quality: 图片质量
    """
    def __init__(self, source: bytes, **kwargs) -> None:
        self._source = source
        self._text = kwargs.get('text') or 'ALOHA'
        self._color = kwargs.get('color') or '#f8f4ed'
        self._font = kwargs.get('font') or os.path.join(
            os.path.dirname(__file__), 'font/hupo.ttf')
        self._space = kwargs.get('space') or 75
        self._opacity = kwargs.get('opacity') or 0.5
        self._angle = kwargs.get('angle') or 30
        self._txtsize = kwargs.get('txtsize') or 50
        self._format = kwargs.get('format') or 'PNG'
        self._height = kwargs.get('height') or 1.3
        self._quality = kwargs.get('quality') or 80
        self._image = Image.open(BytesIO(self._source))
        print(self._text, self._color, self._height)

    @staticmethod
    def set_opacity(im, opacity):
        """设置水印透明度"""
        assert opacity >= 0 and opacity <= 1
        alpha = im.split()[3]
        alpha = ImageEnhance.Brightness(alpha).enhance(opacity)
        im.putalpha(alpha)
        return im

    @staticmethod
    def crop_image(im):
        """剪裁图片边缘"""
        bg = Image.new(mode='RGBA', size=im.size)
        diff = ImageChops.difference(im, bg)
        del bg
        bbox = diff.getbbox()
        if bbox:
            im.crop(bbox)
        return im

    def gen_mark(self):
        """
            生成mark图片，返回添加水印的函数
        """

        # 字体宽度、高度
        width = len(self._text) * self._txtsize
        height = round(self._txtsize * self._height)

        # 创建水印图片(宽度、高度)
        mark = Image.new(mode='RGBA', size=(width, height))

        # 生成文字
        draw_table = ImageDraw.Draw(im=mark)
        draw_table.text(xy=(0, 0),
                        text=self._text,
                        fill=self._color,
                        font=ImageFont.truetype(self._font,
                                                size=self._txtsize))
        del draw_table

        # 裁剪空白
        mark = self.crop_image(mark)

        # 透明度
        self.set_opacity(mark, self._opacity)

        def mark_im(im):
            ''' 在im图片上添加水印 im为打开的原图'''

            # 计算斜边长度
            c = int(
                math.sqrt(im.size[0] * im.size[0] + im.size[1] * im.size[1]))

            # 以斜边长度为宽高创建大图（旋转后大图才足以覆盖原图）
            mark2 = Image.new(mode='RGBA', size=(c, c))

            # 在大图上生成水印文字，此处mark为上面生成的水印图片
            y, idx = 0, 0
            while y < c:
                # 制造x坐标错位
                x = -int((mark.size[0] + self._space) * 0.5 * idx)
                idx = (idx + 1) % 2

                while x < c:
                    # 在该位置粘贴mark水印图片
                    mark2.paste(mark, (x, y))
                    x = x + mark.size[0] + self._space
                y = y + mark.size[1] + self._space

            # 将大图旋转一定角度
            mark2 = mark2.rotate(self._angle)

            # 在原图上添加大图水印
            if im.mode != 'RGBA':
                im = im.convert('RGBA')
            im.paste(
                mark2,  # 大图
                (int((im.size[0] - c) / 2), int((im.size[1] - c) / 2)),  # 坐标
                mask=mark2.split()[3])
            del mark2
            return im

        return mark_im

    def add_mark(self, mark):
        """
            添加水印, 然后保存图片
        """
        self._image = mark(self._image)
        out = BytesIO()
        self._image.save(out, quality=self._quality, format='PNG')
        return out

    def water_mark(self):
        """生成水印图片"""
        mark = self.gen_mark()
        out = self.add_mark(mark)
        return out


if __name__ == '__main__':
    with open('C:/Users/Fisher/Desktop/a.jpg', 'rb') as f:
        a = f.read()
    marker = WaterMark(a, angle=90, text='ALOHA')
    out = marker.water_mark()
    image = Image.open(out)
    image.show()
